package ai

import (
	"context"
	"encoding/json"
	"io"
	"strings"
	"time"

	"paper-pro/app/chat/api/internal/svc"
	"paper-pro/app/chat/api/internal/types"
	"paper-pro/app/chat/api/internal/types/role"
	kqtypes "paper-pro/app/chat/jobs/types"
	"paper-pro/pkg/ctxdata"
	"paper-pro/pkg/rpc/chat/client/chatdialogservice"
	"paper-pro/pkg/rpc/rag/pb/rag"
	"paper-pro/pkg/xerr"

	"github.com/google/uuid"
	"github.com/jinzhu/copier"
	"github.com/pkg/errors"
	"github.com/zeromicro/go-zero/core/logx"
)

var ErrCompletionsCancel = xerr.NewErrCode(xerr.REQUEST_CANCEL)

type CompletionsRecv struct {
	AiReply *types.ChatMsg
	Err     error
}

type CompletionsLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewCompletionsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CompletionsLogic {
	return &CompletionsLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *CompletionsLogic) pushKqChatMsg(msg *types.ChatMsg, userId int64) error {
	kqMsg := kqtypes.KqChatMsg{
		Content:   msg.Content,
		Role:      msg.Role,
		ChatAt:    msg.ChatAt,
		DialogId:  msg.DialogID,
		UserId:    userId,
		RequestId: msg.RequestId,
	}
	msgBytes, err := json.Marshal(kqMsg)
	if err != nil {
		return err
	}
	return l.svcCtx.KqPusherClient.Push(l.ctx, string(msgBytes))
}

func (l *CompletionsLogic) getOrCreateChatDialog(req *types.UserMsg, userId int64) (dialog *types.ChatDialog, err error) {
	// todo: 优先考虑从redis中获取

	dialog = &types.ChatDialog{}
	if req.IsNew {
		resp, err := l.svcCtx.ChatDialogRpc.Create(l.ctx, &chatdialogservice.ChatDialogCreateReq{
			Dialog: &chatdialogservice.ChatDialog{
				Name:         req.ChatDialog.Name,
				LastActiveAt: 0, // 无所谓
				AssistantId:  req.ChatDialog.AssistantId,
				UserId:       userId,
				RequestId:    req.ChatMsg.RequestId,
			},
		})
		if err != nil {
			return nil, errors.Wrapf(err, "CompletionsLogic: getOrCreateChatDialog: ChatDialogRpc Create error: %v", err)
		}
		err = copier.Copy(dialog, resp.Dialog)
		if err != nil {
			return nil, errors.Wrapf(err, "CompletionsLogic: getOrCreateChatDialog: Create Copy(dialog, resp.Dialog) error: %v", err)
		}
		req.ChatMsg.DialogID = dialog.Id // 此处将UserMsg的DialogID赋值
	} else {
		resp, err := l.svcCtx.ChatDialogRpc.Get(l.ctx, &chatdialogservice.ChatDialogGetReq{
			Id:     req.ChatMsg.DialogID,
			UserId: userId,
		})
		if err != nil {
			return nil, errors.Wrapf(err, "CompletionsLogic: getOrCreateChatDialog: ChatDialogRpc Get error: %v", err)
		}
		err = copier.Copy(dialog, resp.Dialog)
		if err != nil {
			return nil, errors.Wrapf(err, "CompletionsLogic: getOrCreateChatDialog: Get Copy(dialog, resp.Dialog) error: %v", err)
		}
	}
	return
}

func (l *CompletionsLogic) Completions(req *types.UserMsg) <-chan CompletionsRecv {
	recvChan := make(chan CompletionsRecv, 4)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				recvChan <- CompletionsRecv{Err: errors.Errorf("CompletionsLogic Completions: Unexpected panic: %v", r)}
			}
		}()

		defer func() {
			close(recvChan)
		}()

		userId, err := ctxdata.GetUidFromCtx(l.ctx)
		if err != nil {
			recvChan <- CompletionsRecv{
				Err: errors.Wrapf(err, "CompletionsLogic failed: get UserId error: %v", err),
			}
			return
		}

		// todo: 测一下时间，可以考虑搞一个map-reduce
		dialog, err := l.getOrCreateChatDialog(req, userId)
		if err != nil {
			recvChan <- CompletionsRecv{
				Err: errors.Wrapf(err, "CompletionsLogic failed: Get Or Create Chat Dialog error: %v", err),
			}
			return
		}

		aiStream, err := l.svcCtx.RagRpc.ChatStream(l.ctx, &rag.UserMsg{
			Content:     req.ChatMsg.Content,
			AssistantId: dialog.AssistantId,
			DialogId:    dialog.Id,
			UserId:      userId,
		})
		if err != nil {
			recvChan <- CompletionsRecv{
				Err: errors.Wrapf(err, "NewCompletionsLogic failed: 生成sse流失败，error: %v", err),
			}
			return
		}
		defer func() {
			aiStream.CloseSend()
		}()

		buffer := strings.Builder{}
		aiMsg := &types.ChatMsg{
			Content:   "",
			Role:      role.Ai.String(),
			ChatAt:    time.Now().Add(time.Second).Unix(), // 因为精度为秒，怕与用户消息区分不开时间
			DialogID:  dialog.Id,
			RequestId: uuid.New().String(),
		}
		for {
			reply, err := aiStream.Recv()
			if err != nil {
				if err == io.EOF {
					// 存储聊天记录，将聊天记录抛给kafka消息队列（忽略错误处理）
					_ = l.pushKqChatMsg(&req.ChatMsg, userId)
					_ = l.pushKqChatMsg(aiMsg, userId)
				} else if errors.Is(err, context.Canceled) {
					err = ErrCompletionsCancel
				} else {
					err = errors.Wrapf(err, "CompletionsLogic Completions: grpc Recv error: %v", err)
				}
				recvChan <- CompletionsRecv{
					Err: err,
				}
				break
			}

			_, err = buffer.WriteString(reply.Chunk)
			if err != nil {
				recvChan <- CompletionsRecv{
					Err: errors.Wrapf(err, "CompletionsLogic Completions: buffer WriteString failed: %v", err),
				}
				break
			}
			aiMsg.Content = buffer.String()
			recvChan <- CompletionsRecv{
				AiReply: aiMsg,
			}
		}
	}()

	return recvChan
}
